#!/usr/bin/env python
# -*- coding: utf-8 -*-
import rospy
import numpy as np
import tf
import math
from tf import transformations
from collections import deque
from nav_msgs.msg import Odometry
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped
from geometry_msgs.msg import Twist

id = "/tb3_0/odom"
index = 0



class Prediction(object):
    def __init__(self, size = 30, prediction_horizons = 100, dt = 0.5, desired_linear_vel = 0.20):
        self.pub_show = rospy.Publisher('show_target', Odometry, queue_size=1)
        self.pub_desired_path = rospy.Publisher('/desired_path', Path, queue_size=1)
        self.sub_cmd = rospy.Subscriber('/cmd_user', Twist, self.cmd_cb, queue_size=1,buff_size=52428800) 

        self.autonomy = 0
        self.traj_buffer = deque(maxlen=size)
        self.traj_len = 0
        self.maxsize = size
        self.prediction_horizons = prediction_horizons
        self.desired_linear_vel = desired_linear_vel
        self.dt = dt

        self.initial_pos = Odometry()
        self.sum_traj_pos = Odometry()
        self.target_pos = Odometry()
        self.pub_path = Path()
    
    def cmd_cb(self, data):
        self.autonomy = data.linear.z
        # print("self.autonomy: %lf" %self.autonomy)

    def run(self, latest_leader_pos):
        global index
        index += 1
        self.update_traj(latest_leader_pos)
        if self.traj_len == self.maxsize:
            self.traj_mean()
            self.desired_path()

    def desired_path(self):
        global id

        x_mean = self.initial_pos.pose.pose.position.x
        y_mean = self.initial_pos.pose.pose.position.y
        a = self.initial_pos.pose.pose.orientation.x
        b = self.initial_pos.pose.pose.orientation.y
        c = self.initial_pos.pose.pose.orientation.z
        d = self.initial_pos.pose.pose.orientation.w
        (roll,pitch,yaw) = transformations.euler_from_quaternion([a, b, c, d])
        theta_mean = yaw
        linear_vel_mean = self.initial_pos.twist.twist.linear.x
        angular_vel_mean = self.initial_pos.twist.twist.angular.z

        x_t_1 = x_mean
        y_t_1 = y_mean
        yaw_t_1 = yaw
        v_t_1 = self.desired_linear_vel
        w_t_1 = angular_vel_mean

        # 预测轨迹
        iter = int(self.prediction_horizons / self.dt)
        desired_path = Path()

        for t in range(0, iter + 1):
            desired_path.header.stamp = rospy.get_rostime()
            desired_path.header.frame_id = id
            desired_path.header.seq = t

            pose = PoseStamped()
            pose.header.seq = t 
            pose.header.frame_id = id
            pose.header.stamp = rospy.get_rostime()
            pose.pose.position.x = x_t_1
            pose.pose.position.y = y_t_1
            theta_target = yaw_t_1
            q = tf.transformations.quaternion_from_euler(0, 0, theta_target)
            pose.pose.orientation.x = q[0]
            pose.pose.orientation.y = q[1]
            pose.pose.orientation.z = q[2]
            pose.pose.orientation.w = q[3]

            x_t_0 = x_t_1
            y_t_0 = y_t_1
            yaw_t_0 = yaw_t_1
            v_t_0 = v_t_1
            w_t_0 = w_t_1

            x_t_1 = x_t_0 + v_t_0 * self.dt * math.cos(yaw_t_0)
            y_t_1 = y_t_0 + v_t_0 * self.dt * math.sin(yaw_t_0)
            yaw_t_1 = yaw_t_0 + w_t_0 * self.dt
            v_t_1 = v_t_0
            w_t_1 = w_t_0

            desired_path.poses.append(pose)
        if abs(self.autonomy - 1.0) >= 0.05:
            self.pub_path = desired_path
        self.pub_desired_path.publish(self.pub_path)

    
    def traj_mean(self):
        # pose
        self.initial_pos.header = self.sum_traj_pos.header
        self.initial_pos.child_frame_id = self.sum_traj_pos.child_frame_id
        self.initial_pos.pose.pose.position.x = self.sum_traj_pos.pose.pose.position.x / self.traj_len
        self.initial_pos.pose.pose.position.y = self.sum_traj_pos.pose.pose.position.y / self.traj_len
        self.initial_pos.pose.pose.position.z = self.sum_traj_pos.pose.pose.position.z / self.traj_len
        self.initial_pos.pose.pose.orientation.x = self.sum_traj_pos.pose.pose.orientation.x / self.traj_len
        self.initial_pos.pose.pose.orientation.y = self.sum_traj_pos.pose.pose.orientation.y / self.traj_len
        self.initial_pos.pose.pose.orientation.z = self.sum_traj_pos.pose.pose.orientation.z / self.traj_len
        self.initial_pos.pose.pose.orientation.w = self.sum_traj_pos.pose.pose.orientation.w / self.traj_len
        # twist
        self.initial_pos.twist.twist.linear.x = self.sum_traj_pos.twist.twist.linear.x / self.traj_len
        self.initial_pos.twist.twist.linear.y = self.sum_traj_pos.twist.twist.linear.y / self.traj_len
        self.initial_pos.twist.twist.linear.z = self.sum_traj_pos.twist.twist.linear.z / self.traj_len
        self.initial_pos.twist.twist.angular.x = self.sum_traj_pos.twist.twist.angular.x / self.traj_len
        self.initial_pos.twist.twist.angular.y = self.sum_traj_pos.twist.twist.angular.y / self.traj_len
        self.initial_pos.twist.twist.angular.z = self.sum_traj_pos.twist.twist.angular.z / self.traj_len

    def update_traj(self, latest_leader_pos):   
        # debug
        # if abs(self.autonomy - 1.0) <= 0.05:
        #     return
        self.traj_len += 1
        self.add_to_point(latest_leader_pos)  
        if self.traj_len > self.maxsize:
            self.traj_len = self.maxsize
            self.cut_to_point(self.traj_buffer[0])
        self.traj_buffer.append(latest_leader_pos)

        # print("traj_buffer_len : %d" %(len(self.traj_buffer)))
        # print("1111111111111111111111111111111111111111111")
        # print(self.traj_buffer)
        # print("2222222222222222222222222222222222222222222")
        # print(self.traj_buffer[0])  # 队尾（最早的数据）
        # print("3333333333333333333333333333333333333333333")
        # print(self.traj_buffer[-1]) # 队首（最新的数据）
        # print("4444444444444444444444444444444444444444444")
        # print("traj_len: %d" %self.traj_len)
    
    def add_to_point(self, latest_leader_pos):
        # pose
        self.sum_traj_pos.header = latest_leader_pos.header
        self.sum_traj_pos.child_frame_id = latest_leader_pos.child_frame_id
        self.sum_traj_pos.pose.pose.position.x += latest_leader_pos.pose.pose.position.x
        self.sum_traj_pos.pose.pose.position.y += latest_leader_pos.pose.pose.position.y
        self.sum_traj_pos.pose.pose.position.z += latest_leader_pos.pose.pose.position.z
        self.sum_traj_pos.pose.pose.orientation.x += latest_leader_pos.pose.pose.orientation.x
        self.sum_traj_pos.pose.pose.orientation.y += latest_leader_pos.pose.pose.orientation.y
        self.sum_traj_pos.pose.pose.orientation.z += latest_leader_pos.pose.pose.orientation.z
        self.sum_traj_pos.pose.pose.orientation.w += latest_leader_pos.pose.pose.orientation.w
        # twist
        self.sum_traj_pos.twist.twist.linear.x += latest_leader_pos.twist.twist.linear.x
        self.sum_traj_pos.twist.twist.linear.y += latest_leader_pos.twist.twist.linear.y
        self.sum_traj_pos.twist.twist.linear.z += latest_leader_pos.twist.twist.linear.z
        self.sum_traj_pos.twist.twist.angular.x += latest_leader_pos.twist.twist.angular.x
        self.sum_traj_pos.twist.twist.angular.y += latest_leader_pos.twist.twist.angular.y
        self.sum_traj_pos.twist.twist.angular.z += latest_leader_pos.twist.twist.angular.z

    def cut_to_point(self, first_pos):
        # pose
        self.sum_traj_pos.pose.pose.position.x -= first_pos.pose.pose.position.x
        self.sum_traj_pos.pose.pose.position.y -= first_pos.pose.pose.position.y
        self.sum_traj_pos.pose.pose.position.z -= first_pos.pose.pose.position.z
        self.sum_traj_pos.pose.pose.orientation.x -= first_pos.pose.pose.orientation.x
        self.sum_traj_pos.pose.pose.orientation.y -= first_pos.pose.pose.orientation.y
        self.sum_traj_pos.pose.pose.orientation.z -= first_pos.pose.pose.orientation.z
        self.sum_traj_pos.pose.pose.orientation.w -= first_pos.pose.pose.orientation.w
        # twist
        self.sum_traj_pos.twist.twist.linear.x -= first_pos.twist.twist.linear.x
        self.sum_traj_pos.twist.twist.linear.y -= first_pos.twist.twist.linear.y
        self.sum_traj_pos.twist.twist.linear.z -= first_pos.twist.twist.linear.z
        self.sum_traj_pos.twist.twist.angular.x -= first_pos.twist.twist.angular.x
        self.sum_traj_pos.twist.twist.angular.y -= first_pos.twist.twist.angular.y
        self.sum_traj_pos.twist.twist.angular.z -= first_pos.twist.twist.angular.z


def leaderOdom_CB(data):
    global latest_leader_pos
    latest_leader_pos = data
    
def node():
    global latest_leader_pos
    rospy.init_node('prediction_node')
    rospy.loginfo("prediction_node is started!!")
    prediction_module = Prediction()

    rospy.Subscriber('/tb3_0/odom', Odometry, leaderOdom_CB, queue_size=1,buff_size=52428800)    # 接收领航者位置速度信息
    rate = rospy.Rate(10) # 10hz
    while not rospy.is_shutdown():
        print("x: %lf, y: %lf, v: %lf, w: %lf"%(latest_leader_pos.pose.pose.position.x, latest_leader_pos.pose.pose.position.y,\
                                                latest_leader_pos.twist.twist.linear.x, latest_leader_pos.twist.twist.angular.z))
        prediction_module.run(latest_leader_pos)
        rate.sleep()


if __name__ == '__main__':
    latest_leader_pos = Odometry()
    try:
        node()
    except rospy.ROSInterruptException:
        pass    
